import argparse
import torch
import random
import numpy as np
import os


def get_params():

    parser = argparse.ArgumentParser()

    parser.add_argument('--seed', type=int, default=369, help='Random seed.')

    # parameters for gcl learning
    parser.add_argument('--epochs_gcl', type=int, default=100, help='Number of epochs to train gcl model.')
    parser.add_argument('--gcl_lr', type=float, default=0.01, help='Initial learning rate for gcl learning.')
    parser.add_argument('--input_dim_encoder', type=int, default=64, help='Number of input units.')
    parser.add_argument('--hidden_dim_encoder', type=int, default=32, help='Number of hidden units.')
    parser.add_argument('--num_layers_encoder', type=int, default=2, help='Number of layers for gnn encoder.')


    # parameters for dataset
    # data infor
    parser.add_argument('--train_ratio', type=float, default=0.8,
                        help='train_ratio.')
    parser.add_argument('--data_name', type=str,
                        help='may be crg_gnp_random_graph; rpt_rt_tree_graph; rc_bg_graph')
    parser.add_argument('--graph_name', type=str,default='crg_gnp_0.2',
                        help='may be crg_gnp_p, p=0.2~0.9; rpt_rt; rc_bg')
    parser.add_argument('--label_list', type=list, default=['crg', 'gnp'],
                        help="may be ['crg', 'gnp'], ['rpt', 'rt'], ['rc', 'bg']")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument('--device', type=str, default=device,
                        help='device to use.')
    parser.add_argument('--exp_name', type=str, default='salmas',
                        help='type of experiment.')
    parser.add_argument('--step', type=str, default='gcl',
                        help='type of learning or solving.')
    parser.add_argument('--dataset', type=str, default='salmas_data_1',
                        help='total dataset.')


    args = parser.parse_args()

    return args


def set_seed(seed):
    """Set seed"""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.enable = True
    os.environ["PYTHONHASHSEED"] = str(seed)


args = get_params()
# os.environ["CUDA_VISIBLE_DEVICES"] = '0'
set_seed(args.seed)

